package com.dragon.datastructure.tree.trie;

import java.util.TreeMap;

/**
 * trie
 *
 * @author dragon
 * @since 2020/1/2
 */
public class Trie {

    private int size;

    private Node root;

    private static class Node {
        boolean isWord;
        TreeMap<Character, Node> next;

        public Node(boolean isWord, TreeMap<Character, Node> next) {
            this.isWord = isWord;
            this.next = next;
        }

        public Node() {
            next = new TreeMap<>();
        }
    }

    public Trie() {
        root = new Node();
    }

    public void add(String word) {
        if (word == null || word.trim().length() == 0) {
            return;
        }
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (cur.next.get(c) == null) {
                cur.next.put(c, new Node());
            }
            cur = cur.next.get(c);
        }

        if (!cur.isWord) {
            cur.isWord = true;
            size++;
        }
    }

    public void addCur(String word) {
        if (word == null || word.trim().length() == 0) {
            return;
        }
        doAddCur(root, word, 0);
    }

    private void doAddCur(Node cur, String word, int index) {
        if (word.length() == index) {
            if (!cur.isWord) {
                cur.isWord = true;
                size++;
            }
            return;
        }
        char c = word.charAt(index);
        if (cur.next.get(c) == null) {
            cur.next.put(c, new Node());
        }
        doAddCur(cur.next.get(c), word, index + 1);
    }

    public boolean contains(String word) {
        if (word == null || word.trim().length() == 0) {
            return false;
        }
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }

        return cur.isWord;
    }


    public boolean isPrefix(String prefix) {
        if (prefix == null || prefix.trim().length() == 0) {
            return false;
        }
        Node cur = root;
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (cur.next.get(c) == null) {
                return false;
            }

            cur = cur.next.get(c);
        }

        return true;
    }
}
